package org.lanyu.springainovel.common.service.impl;

import lombok.extern.slf4j.Slf4j;
import org.lanyu.springainovel.common.dto.StreamGenerationRequest;
import org.lanyu.springainovel.common.service.StreamGenerationService;
import org.lanyu.springainovel.common.config.ChatModelManager;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 通用流式生成服务实现
 */
@Slf4j
@Service
public class StreamGenerationServiceImpl implements StreamGenerationService {
    
    @Autowired
    private ChatModelManager chatModelManager;
    
    @Autowired
    private org.lanyu.springainovel.chat.service.PromptTemplateService promptTemplateService;
    
    // 存储活跃的流式生成任务，用于终止操作
    private static final ConcurrentHashMap<String, AtomicBoolean> activeStreams = new ConcurrentHashMap<>();
    
    @Override
    public Flux<String> generateStreamResponse(StreamGenerationRequest request) {
        // 为每个流式生成任务创建一个终止信号
        String abortSignal = request.getAbortSignal();
        if (abortSignal != null && !abortSignal.isEmpty()) {
            AtomicBoolean isAborted = new AtomicBoolean(false);
            activeStreams.put(abortSignal, isAborted);
            
            log.info("开始流式生成任务，终止信号: {}", abortSignal);
        }
        
        try {
            // 1. 获取AI模型
            ChatModel chatModel = chatModelManager.getOrCreateChatModelById(request.getModelId());
            if (chatModel == null) {
                return Flux.just("无法创建聊天模型");
            }
            
            // 2. 获取系统提示词
            String systemPrompt = getSystemPrompt(request.getSystemPromptId());
            if (systemPrompt == null) {
                return Flux.just("无法获取系统提示词模板");
            }
            
            // 3. 创建ChatClient
            ChatClient chatClient = ChatClient.builder(chatModel).build();
            
            // 4. 构建错误消息前缀
            String errorPrefix = request.getBusiness() != null ? 
                request.getBusiness() + "失败" : "生成失败";
            
            // 5. 生成流式响应
            Flux<String> flux = chatClient.prompt()
                    .system(systemPrompt)
                    .user(request.getUserPrompt())
                    .stream()
                    .content()
                    .doOnError(error -> log.error("{}流式响应失败", errorPrefix, error))
                    .onErrorResume(error -> Flux.just(errorPrefix + ": " + 
                        (error.getMessage() != null ? error.getMessage() : "未知错误")));
            
            // 如果有终止信号，添加终止检查和清理逻辑
            if (abortSignal != null && !abortSignal.isEmpty()) {
                AtomicBoolean isAborted = activeStreams.get(abortSignal);
                if (isAborted != null) {
                    flux = flux.takeWhile(data -> {
                        if (isAborted.get()) {
                            log.info("流式生成任务已被终止，终止信号: {}", abortSignal);
                            return false;
                        }
                        return true;
                    }).doFinally(signalType -> {
                        // 流结束时清理终止信号
                        activeStreams.remove(abortSignal);
                        log.info("流式生成任务结束，终止信号: {}", abortSignal);
                    });
                }
            }
            
            return flux;
        } catch (Exception e) {
            // 清理终止信号
            if (abortSignal != null && !abortSignal.isEmpty()) {
                activeStreams.remove(abortSignal);
            }
            
            String errorPrefix = request.getBusiness() != null ? 
                request.getBusiness() + "失败" : "生成失败";
            log.error(errorPrefix, e);
            return Flux.just(errorPrefix + ": " + e.getMessage());
        }
    }
    
    @Override
    public boolean abortStreamGeneration(String abortSignal) {
        if (abortSignal == null || abortSignal.isEmpty()) {
            log.warn("终止信号为空，无法终止流式生成任务");
            return false;
        }
        
        AtomicBoolean isAborted = activeStreams.get(abortSignal);
        if (isAborted != null) {
            isAborted.set(true);
            log.info("已设置终止信号，流式生成任务将被终止: {}", abortSignal);
            return true;
        } else {
            log.warn("未找到对应的流式生成任务，终止信号: {}", abortSignal);
            return false;
        }
    }

    /**
     * 获取系统提示词
     * 
     * @param systemPromptId 系统提示词ID
     * @return 系统提示词
     */
    private String getSystemPrompt(Long systemPromptId) {
        if (systemPromptId == null) {
            return null;
        }
        
        // 从服务获取提示词模板
        org.lanyu.springainovel.common.entity.MyPromptTemplate promptTemplate = 
            promptTemplateService.getTemplateById(systemPromptId);
        if (promptTemplate == null || promptTemplate.getContent() == null || 
            promptTemplate.getContent().isEmpty()) {
            return null;
        }
        
        String systemPrompt = promptTemplate.getContent();
        
        // 增加模板使用次数
        promptTemplateService.incrementUsageCount(systemPromptId);
        
        return systemPrompt;
    }
}